import numpy as np
from pymoo.algorithms.moo.nsga2 import NSGA2
from pymoo.core.problem import Problem
from pymoo.operators.crossover.sbx import SBX
from pymoo.operators.mutation.pm import PM
from pymoo.operators.sampling.rnd import FloatRandomSampling
from pymoo.optimize import minimize
from pymoo.visualization.scatter import Scatter
from scipy.constants import c

import obj

# HALF = {
#     'h': 800,
#     'energy': 2.2e9,  # ev
#     'circumference': 479.86,  # m
#     'alpha_c': 9.4e-5,
#     'Q': 1e5,  # 欧姆
#     'RoverQ': 45,  # 欧姆
#     'v_mc': 1.2e6,  # V
#     'phi_s': 158 / 180 * np.pi,  # rad
#     'omega_r': 499.98e6 * 2 * np.pi,  # Hz
# }

# HALF = {
#     'h': 800,
#     'energy': 2.2e9,  # ev
#     'circumference': 479.86,  # m
#     'alpha_c': 9.4e-5,
#     'Q': 1e5,  # 欧姆
#     'RoverQ': 90,  # 欧姆
#     'v_mc': 1.2e6,  # V
#     # 'phi_s': 158 / 180 * np.pi,  # rad
#     'phi_s': 2.780020615501037, # rad
#     'omega_r': (3 * 499.98e6+120e3) * 2 * np.pi,  # Hz
# }
# HALF['omega_r'] = (3 * HALF['h'] * c / HALF['circumference'] + 120e3) * 2 * np.pi
#
# omega_s = obj.omega_s(h=HALF['h'], omega_0=2 * np.pi * HALF['circumference'] / c, alpha_c=HALF['alpha_c'],
#                       v_mc=HALF['v_mc'], phi_s=HALF['phi_s'], energy=HALF['energy'], e=1)
#
# # tau = np.array([200,500,1000]) * 1e-9
# tau = 2.000800834022316e-09
# tau = 100*tau
#
# p = 3
# # p = np.array([i for i in range(-p, p + 1)])

HALF = {
    'h': 800,
    'energy': 2.2e9,  # ev
    'circumference': 479.86,  # m
    'alpha_c': 9.4e-5,
    'Q': 1e5,  # 欧姆
    'RoverQ': 90,  # 欧姆
    'v_mc': 1.2e6,  # V
    # 'phi_s': 158 / 180 * np.pi,  # rad
    'phi_s': 2.780020615501037,  # rad
    'omega_r': (3 * 499.98e6 + 120e3) * 2 * np.pi,  # Hz
}
HALF['omega_r'] = (3 * HALF['h'] * c / HALF['circumference'] + 120e3) * 2 * np.pi
omega_s = obj.omega_s(h=HALF['h'], omega_0=2 * np.pi * c / HALF['circumference'], alpha_c=HALF['alpha_c'],
                      v_mc=HALF['v_mc'], phi_s=HALF['phi_s'], energy=HALF['energy'], e=1)
tau = 2.000800834022316e-09
tau = 100*tau
p = 3

class MyMultiObjectiveProblem(Problem):
    def __init__(self):
        # 定义2个变量，范围在[-5, 5]
        super().__init__(n_var=5,
                         n_obj=2,  # 2个目标函数
                         n_constr=2,  # 无约束
                         xl=[0, -0.3 * np.pi, 0.15, 0.5, -1,],
                         xu=[12, 0.3 * np.pi, 0.21, 0.8, 1,])

    def _evaluate(self, x, out, *args, **kwargs):
        my_obi = obj.MyObj(omega_0=2 * np.pi * c / HALF['circumference'],
                           omega_complex=omega_s,
                           r_l=HALF['RoverQ'] * HALF['Q'],
                           q_l=HALF['Q'],
                           omega_r=HALF['omega_r'],
                           tau=tau,
                           k_p=x[:, 0], phi=x[:, 1], g=x[:, 2], k=x[:, 3], t_g=np.array([int(i)/100 for i in x[:, 4]*100]),
                           t_0=HALF['circumference'] / c)
        # 第一个目标函数：最小化 x^2 + y^2
        # f1 = x[0] ** 2 + x[1] ** 2
        f1 = my_obi.p_re(p=p, m=800, mu=1)
        f2 = my_obi.p_re(p=p, m=800, mu=2)

        f_0 = 0.7*f1+0.3*f2

        # 第二个目标函数：最小化 (x-2)^2 + (y-2)^2
        # f2 = (x[0] - 2) ** 2 + (x[1] - 2) ** 2
        f3 = my_obi.p_im(p=p, m=800, mu=1)

        # out["F"] = [f1, f2, f3]
        out["F"] = [f_0, f3]

        # # 注意：pymoo中约束格式为 g(x) ≤ 0
        omega = 1000000 * np.linspace(-np.pi, np.pi, 1000) + HALF['omega_r']
        # zeros = np.vectorize(my_obi.phase_margin)(omega)
        f100,f101 = my_obi.phase_margin(omega,-150,150)

        out['G'] = [f100,f101]


# 创建问题实例
problem = MyMultiObjectiveProblem()

algorithm = NSGA2(
    pop_size=500,  # 种群大小
    sampling=FloatRandomSampling(),  # 随机采样
    crossover=SBX(prob=0.9, eta=15),  # 模拟二进制交叉
    mutation=PM(prob=0.1, eta=20),  # 多项式变异
    eliminate_duplicates=True  # 消除重复个体
)

res = minimize(problem,
               algorithm,
               ('n_gen', 500),  # 迭代200代
               verbose=True,
               seed=1)

print("完成!")
print(f"找到的Pareto前沿解数量: {len(res.X)}")

plot = Scatter(title="Pareto")
plot.add(res.F, color="red")
plot.show(block=True)

idx = (res.F[:,0] == min(res.F[:,0]))
idx = [i for i in range(len(idx)) if idx[i]]

# for i,resi in enumerate(res.F[:,0]):
#     if resi < 1e7 && resi < 1e7:
#         print (i)
#
# for i,resi in enumerate(res.F[:,1]):
#     if resi < 1e7:
#         print (i)

# target = np.sum(res.X,axis=1)

print(f'idx: {idx[0]}, kp: {res.X[idx,0][0]}, phi_s/pi: {res.X[idx,1][0]/np.pi}, G: {res.X[idx,2][0]}, K: {res.X[idx,3][0]}, T_g/T_0: {res.X[idx,4][0]:.2f}')


# my_obi = obj.MyObj(omega_0=2 * np.pi * c / HALF['circumference'],
#                    omega_complex=omega_s,
#                    r_l=HALF['RoverQ'] * HALF['Q'],
#                    q_l=HALF['Q'],
#                    omega_r=HALF['omega_r'],
#                    tau=tau,
#                    k_p=res.X[idx,0][0],
#                    phi=res.X[idx,1][0], g=res.X[idx,2][0], k=res.X[idx,3][0], t_g=res.X[idx,4][0],
#                    t_0=HALF['circumference'] / c)
# print(f"p_re1 = {my_obi.p_re(p=p, m=800, mu=1):.5e}")
# print(f"p_re2 = {my_obi.p_re(p=p, m=800, mu=2):.5e}")
# print(f"p_im1 = {my_obi.p_im(p=p, m=800, mu=1):.5e}")
# print(f"T_g/T_0: {res.X[idx,4][0]}")

my_obj = obj.MyObj(omega_0=2 * np.pi * c / HALF['circumference'],
                   omega_complex=omega_s,
                   r_l=HALF['RoverQ'] * HALF['Q'],
                   q_l=HALF['Q'],
                   omega_r=HALF['omega_r'],
                   tau=tau,
                   k_p=res.X[idx,0][0],
                   phi=res.X[idx,1][0], g=res.X[idx,2][0], k=res.X[idx,3][0], t_g=int(res.X[idx,4][0]*100)/100,
                   t_0=HALF['circumference'] / c)
print(f"p_re1 = {my_obj.p_re(p=p, m=800, mu=1):.5e}")
print(f"p_re2 = {my_obj.p_re(p=p, m=800, mu=2):.5e}")
print(f"p_im1 = {my_obj.p_im(p=p, m=800, mu=1):.5e}")
print(f"该点相位裕度：{my_obj.phase_margin(1000000 * np.linspace(-np.pi, np.pi, 1000) + HALF['omega_r'],-150,150)}")

# my_obi = obj.MyObj(omega_0=2 * np.pi * c / HALF['circumference'],
#                    omega_complex=omega_s,
#                    r_l=HALF['RoverQ'] * HALF['Q'],
#                    q_l=HALF['Q'],
#                    omega_r=HALF['omega_r'],
#                    tau=tau,
#                    k_p=res.X[idx,0][0],
#                    phi=res.X[idx,1][0], g=res.X[idx,2][0], k=res.X[idx,3][0], t_g=-0.81,
#                    t_0=HALF['circumference'] / c)
# print(f"p_re1 = {my_obi.p_re(p=p, m=800, mu=1):.5e}")
# print(f"p_re2 = {my_obi.p_re(p=p, m=800, mu=2):.5e}")
# print(f"p_im1 = {my_obi.p_im(p=p, m=800, mu=1):.5e}")